# - * - coding : utf - 8 - * -
#
# Copyright(c) 2025 Huawei Technologies                                            Co., Ltd.
# openUBMC is licensed under Mulan PSL                                             v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2                                            at:
# http: // license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON - INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more           details.

"""
await 表达式检测工具

检测表达式树中是否包含 await 节点，以及提取所有 await 节点。
"""

from typing import List
from mcpy.ir.base import IRNode
from mcpy.ir.expressions.await_expr import IRAwait


def contains_await(node: IRNode) -> bool:
    """检测节点树中是否包含 await 表达式"""
    if isinstance(node, IRAwait):
        return True

    for child in node.get_child_nodes():
        if contains_await(child):
            return True

    return False


def collect_await_nodes(node: IRNode) -> List[IRAwait]:
    """收集节点树中的所有 await 表达式"""
    result = []

    if isinstance(node, IRAwait):
        result.append(node)

    # 递归收集所有子节点中的 await
    for child in node.get_child_nodes():
        result.extend(collect_await_nodes(child))

    return result


def count_await_nodes(node: IRNode) -> int:
    """统计节点树中 await 表达式的数量"""
    return len(collect_await_nodes(node))
